# -*- coding: UTF-8 -*-
from math import log, log2

# 第一部分
"""
Parameters:
    无
Returns:
    dataSet - 数据集
    labels - 分类属性
"""
def createDataSet():
    dataSet = [[0, 0, 0, 0, 'no'],         #数据集
            [0, 0, 0, 1, 'no'],
            [0, 1, 0, 1, 'yes'],
            [0, 1, 1, 0, 'yes'],
            [0, 0, 0, 0, 'no'],
            [1, 0, 0, 0, 'no'],
            [1, 0, 0, 1, 'no'],
            [1, 1, 1, 1, 'yes'],
            [1, 0, 1, 2, 'yes'],
            [1, 0, 1, 2, 'yes'],
            [2, 0, 1, 2, 'yes'],
            [2, 0, 1, 1, 'yes'],
            [2, 1, 0, 1, 'yes'],
            [2, 1, 0, 2, 'yes'],
            [2, 0, 0, 0, 'no']]
    labels = ['不放贷', '放贷']             #分类属性
    return dataSet, labels                #返回数据集和分类属性
"""
函数说明:计算给定数据集的经验熵(香农熵)
Parameters:
    dataSet - 数据集
Returns:
    shannonEnt - 经验熵(香农熵)
"""
def calcShannonEnt(dataSet):
    numEntires = len(dataSet)                        #返回数据集的行数
    labelCounts = {}                                #保存每个标签(Label)出现次数的字典
    #在下面添加代码，对每组特征向量进行统计，得到不同label的计数，并保存到labelCounts中
    for elem in dataSet:
        label = elem[-1]
        if label not in labelCounts:
            labelCounts[label] = 1
        else:
            labelCounts[label] = labelCounts[label] + 1
    
    shannonEnt = 0.0                                #经验熵(香农熵)
    #在下面添加代码，计算当前数据集基于labelCounts的经验熵，结果存放在shannonEnt中
    for key in labelCounts.keys():
        p = labelCounts[key] / numEntires#求比率
        shannonEnt -= p * log2(p)#H(D)=-∑p*log2(p)

    return shannonEnt                                #返回经验熵(香农熵)

# if __name__ == '__main__':
#     dataSet, features = createDataSet()
#     print(dataSet)
#     print(calcShannonEnt(dataSet))

# 第二部分
# -*- coding: UTF-8 -*-

"""
函数说明:按照给定特征划分数据集，得到第axis维特征值为value的划分数据子集
Parameters:
    dataSet - 待划分的数据集
    axis - 划分数据集的特征
    value - 需要返回的特征的值
Returns:
    无
"""
def splitDataSet(dataSet, axis, value):       
    retDataSet = []             #创建返回的数据集列表
    #在下面添加代码，将dataSet中第axis维特征等于value的数据取出，存放入retDataSet
    #注意，为方便后面的计算，存放入retDataSet的每一组记录需要去掉axis特征
    for elem in dataSet:
        if elem[axis] == value:
            tmp = elem.copy()
            tmp.pop(axis)
            retDataSet.append(tmp)
            
    return retDataSet     #返回划分后的数据集
"""
函数说明:选择最优特征
Parameters:
    dataSet - 数据集
Returns:
    bestFeature - 信息增益最大的(最优)特征的索引值
"""
def chooseBestFeatureToSplit(dataSet):
    numFeatures = len(dataSet[0]) - 1                    #特征数量
    baseEntropy = calcShannonEnt(dataSet)                 #计算数据集的香农熵
    bestInfoGain = 0.0                                  #信息增益
    bestFeature = -1                                    #最优特征的索引值
    for i in range(numFeatures):                         #遍历所有特征
        #获取dataSet的第i个所有特征
        featList = [example[i] for example in dataSet]
        uniqueVals = set(featList)                         #创建set集合{},元素不可重复
        newEntropy = 0.0                                  #经验条件熵
        #在下方添加代码，计算根据uniqueVals中的特征值划分的数据子集的条件熵，for循环中需要调用splitDataSet函数，结果记录到newEntropy中
        for elem in uniqueVals:
            newEntropy += calcShannonEnt(splitDataSet(dataSet, i, elem))#累加dataSet中第i个属性为elem的集合的香农熵
        
        infoGain = baseEntropy - newEntropy                     #信息增益
        print("第%d个特征的增益为%.3f" % (i, infoGain))            #打印每个特征的信息增益
        if (infoGain > bestInfoGain):                             #计算信息增益
            bestInfoGain = infoGain                             #更新信息增益，找到最大的信息增益
            bestFeature = i                                     #记录信息增益最大的特征的索引值
    return bestFeature                                             #返回信息增益最大的特征的索引值

if __name__ == '__main__':
    dataSet, features = createDataSet()
    print("最优特征索引值:" + str(chooseBestFeatureToSplit(dataSet)))